package main

import (
	"fmt"
	"github.com/gorilla/websocket"
	"log"
	"net"
	"net/http"
	"sync"
	"time"
)

// Session 结构体，用于存储每个连接的会话信息
type Session struct {
	ID   string
	Conn *websocket.Conn
}

func TestWebSocket2(port string, isWss bool) {
	server := NewWebSocketServer()
	err := server.Start(":"+port, isWss)
	if err != nil {
		log.Fatal(err)
	}
}

/*func NewWebSocketServer() *WebSocketServer {
	mux := http.NewServeMux()
	return &WebSocketServer{
		mux:    mux,
		handle: WebSocketHandler,
	}
}*/
func NewWebSocketServer() *WebSocketServer {
	return &WebSocketServer{
		mux:      http.NewServeMux(),
		sessions: make(map[string]*Session),
		upgrader: &websocket.Upgrader{
			//ReadBufferSize:  libs.GetValidValue(nwcontext.ReadBufferSize, 1024*2),
			//WriteBufferSize: libs.GetValidValue(nwcontext.WriteBufferSize, 1024*4),
			CheckOrigin: func(r *http.Request) bool { return true }, // disable check
		},
	}
}

/*type WebSocketServer struct {
	mux    *http.ServeMux
	handle func(w http.ResponseWriter, r *http.Request)
}*/
type WebSocketServer struct {
	mux      *http.ServeMux
	upgrader *websocket.Upgrader
	sessions map[string]*Session
	lock     sync.Mutex
}

/*func (s *WebSocketServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	s.mux.ServeHTTP(w, r)
}*/
func (s *WebSocketServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path == "/" {
		s.WebSocketHandler(w, r)
		//WebSocketHandler(w, r)
	} else {
		s.mux.ServeHTTP(w, r)
	}
}

func (s *WebSocketServer) Start(addr string, isWss bool) error {

	//if isWss {
	//	s.mux.Handle("/", s)
	//	return http.ListenAndServeTLS(addr, "./TestItem/TestWebSocketServer3/server.crt", "./TestItem/TestWebSocketServer3/server.key", s)
	//}
	// 创建TCP监听器
	l, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}
	defer l.Close()

	// 注册WebSocket处理函数
	//s.mux.Handle("/", http.HandlerFunc(s.handle))
	s.mux.Handle("/", s)

	// 启动服务器
	log.Printf("Starting WebSocket server on %s", addr)

	if isWss {
		return http.ServeTLS(l, s, "./TestItem/TestWebSocketServer3/server.crt", "./TestItem/TestWebSocketServer3/server.key")
	}
	return http.Serve(l, s)
}

// WebSocketHandler 方法，处理 WebSocket 连接
func (s *WebSocketServer) WebSocketHandler(w http.ResponseWriter, r *http.Request) {
	conn, err := s.upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println("Error during connection upgrade:", err)
		return
	}

	// 创建新的会话
	sessionID := generateSessionID() // 生成唯一的会话 ID
	session := &Session{ID: sessionID, Conn: conn}

	// 将会话添加到管理器
	s.lock.Lock()
	s.sessions[sessionID] = session
	s.lock.Unlock()

	log.Printf("New connection: %s", sessionID)

	defer func() {
		// 连接断开时清理会话
		log.Printf("Client disconnected: %s", sessionID)
		s.lock.Lock()
		delete(s.sessions, sessionID)
		s.lock.Unlock()
		conn.Close()
	}()

	for {
		_, msg, err := conn.ReadMessage()
		if err != nil {
			// 优化错误处理
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) {
				log.Printf("Unexpected close error for %s: %v", sessionID, err)
			}
			break // 发生错误时退出循环
		}
		log.Printf("Received from %s: %s\n", sessionID, msg)

		// 这里可以添加处理消息的逻辑
	}
}

// 生成唯一的会话 ID（可以根据需要实现更复杂的逻辑）
func generateSessionID() string {
	return fmt.Sprintf("%d", time.Now().UnixNano())
}
